import pandas as pd

from services.DayKlineService import *
pd.set_option('display.max_columns', None)  # 展示所有列

# 读取贵州茅台股票历史交易数据
def get_data_no_index(code):
    kline = DayKlineService()
    newdata = kline.getAllData(code)
    newdata = pd.DataFrame(newdata)
    # 新数据准备
    newdata['Close'] = newdata['close'].astype(float)
    newdata['Open'] = newdata['open'].astype(float)
    newdata['High'] = newdata['high'].astype(float)
    newdata['Low'] = newdata['low'].astype(float)
    newdata['Volume'] = newdata['volume'].astype(float)
    newdata['Date'] = newdata['trade_date']
    """
    stock_data = pd.read_csv('d:\\csv_shares_day\\' + str(code) + '-19900101-20211231.csv', encoding="gbk")
    stock_data.columns = ["Date", "code", "name", "Close", "High", "Low", "Open", "pre_close", "subval",
                          "subrate",
                          "turnrate",
                          "Volume", "Money", "all_earn", "exchage_earn"]
    stock_data['trade_date'] = pd.to_datetime(stock_data['Date'])
    stock_data.set_index('trade_date', inplace=True)
    data_daily = stock_data.sort_index(ascending=True)
    dataframe = data_daily[(data_daily['Date'] >= "2019-06-01") & (data_daily['Date'] <= "2019-12-31")]
    newdata = dataframe[["Date","Close","High","Low","Open","Volume"]] 
    """

    return newdata

code = "SH600519"
df = get_data_no_index(code)
print(df)

missing_values = df.isnull().sum()


# print("缺失值数量：")
# print(missing_values)

# 计算RSI指标
def calculate_rsi(data, window=14):
    delta = data['Close'].diff()
    gain = delta.copy()
    loss = delta.copy()
    gain[gain < 0] = 0
    loss[loss > 0] = 0
    avg_gain = gain.rolling(window).mean()
    avg_loss = abs(loss.rolling(window).mean())
    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))
    return rsi


# 调用calculate_rsi函数计算RSI指标
df['RSI'] = calculate_rsi(df)
# print(df)

#  交易信号生成
df['Signal'] = 0
df.loc[df['RSI'] > 70, 'Signal'] = -1
df.loc[df['RSI'] < 30, 'Signal'] = 1
# 打印df对象
# print(df)

### 绘制RSI指标曲线
import matplotlib.pyplot as plt

plt.rcParams['font.family'] = ['SimHei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 设置负号显示
rsi = calculate_rsi(df)  # 计算RSI指标
plt.figure(figsize=(12, 6))
plt.plot(df.index, rsi, label='RSI')
plt.title('RSI指标')
plt.xlabel('日期')
plt.ylabel('RSI')
plt.legend()
plt.grid(True)
plt.show()

###  绘制K线图
import mplfinance as mpf

plt.rcParams['font.family'] = ['SimHei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 设置负号显示
# 重新加载数据
df = get_data_no_index(code)

# 创建日期索引
df['Date'] = pd.to_datetime(df['Date'])
df.set_index('Date', inplace=True)

market_colors = mpf.make_marketcolors(up='red', down='green')
my_style = mpf.make_mpf_style(marketcolors=market_colors)
# 绘制K线图
mpf.plot(df, type='candle',
         figsize=(10, 6),
         mav=(10, 20),
         volume=True,
         style=my_style)

### 绘制价格和交易信号图表
plt.rcParams['font.family'] = ['SimHei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 设置负号显示

# 读取贵州茅台股票历史交易数据
df = get_data_no_index(code)

# 创建日期索引
df['Date'] = pd.to_datetime(df['Date'])
df.set_index('Date', inplace=True)


# 计算RSI指标
def calculate_rsi(data, window=14):
    delta = data['Close'].diff()
    gain = delta.copy()
    loss = delta.copy()
    gain[gain < 0] = 0
    loss[loss > 0] = 0
    avg_gain = gain.rolling(window).mean()
    avg_loss = abs(loss.rolling(window).mean())
    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))
    return rsi


# 计算RSI指标
df['RSI'] = calculate_rsi(df)
#  交易信号生成
df['Signal'] = 0
df.loc[df['RSI'] > 70, 'Signal'] = -1
df.loc[df['RSI'] < 30, 'Signal'] = 1

# 绘制价格和交易信号图表
plt.figure(figsize=(12, 6))
plt.plot(df.index, df['Close'], label='Close Price')
plt.scatter(df[df['Signal'] == 1].index, df[df['Signal'] == 1]['Close'], color='green', marker='^', label='Buy Signal')
plt.scatter(df[df['Signal'] == -1].index, df[df['Signal'] == -1]['Close'], color='red', marker='v', label='Sell Signal')
plt.title('贵州茅台股票价格和交易信号')
plt.xlabel('日期')
plt.ylabel('股价')
plt.legend()
plt.grid(True)
plt.show()
